#!/usr/bin/env python3
"""
工作流API客户端
用于在知识库操作完成后调用工作流API
"""

import requests
import json
from typing import Dict, Any, Optional
from loguru import logger
from config import settings


class WorkflowClient:
    """多工作流API客户端 - 支持不同知识库对应不同工作流"""
    
    def __init__(self):
        pass
    
    def _get_workflow_config_by_dataset(self, dataset_id: str) -> Optional[Dict[str, Any]]:
        """根据知识库ID获取工作流配置"""
        kb_config = settings.get_knowledge_base_config(dataset_id)
        if kb_config:
            return kb_config.get("workflow_config")
        return None
    
    def _get_workflow_config_by_key(self, workflow_key: str) -> Optional[Dict[str, Any]]:
        """根据工作流key获取配置"""
        return settings.get_workflow_config_by_key(workflow_key)
    
    def _make_request_to_workflow(self, workflow_config: Dict[str, Any], 
                                  data: Dict[str, Any]) -> Optional[Dict[str, Any]]:
        """向指定工作流发送请求"""
        if not workflow_config or not workflow_config.get('enabled'):
            logger.debug(f"知识库 {data.get('data', {}).get('dataset', {}).get('id')} 的工作流未启用")
            return None
            
        base_url = workflow_config.get('base_url', '').rstrip('/')
        api_key = workflow_config.get('api_key', '')
        endpoint = workflow_config.get('endpoint', '/api/workflow/kb-update')
        timeout = workflow_config.get('timeout', 30)
        
        if not base_url:
            logger.warning(f"知识库 {data.get('data', {}).get('dataset', {}).get('id')} 的工作流URL未配置")
            return None
        
        # 设置请求头
        headers = {
            'Content-Type': 'application/json',
            'User-Agent': 'Dify-KB-Client/1.0.0'
        }
        
        # 如果配置了API密钥，添加到请求头
        if api_key:
            headers['Authorization'] = f'Bearer {api_key}'
            headers['X-API-Key'] = api_key
        
        url = f"{base_url}{endpoint}"
        
        try:
            logger.info(f"调用工作流API: POST {url}")
            
            response = requests.post(
                url=url,
                headers=headers,
                json=data,
                timeout=timeout
            )
            
            # 检查响应状态
            if response.status_code in [200, 201]:
                result = response.json()
                logger.success(f"工作流API调用成功: {response.status_code}")
                return result
            else:
                logger.error(f"工作流API调用失败: {response.status_code} - {response.text}")
                return None
                
        except requests.exceptions.Timeout:
            logger.error(f"工作流API调用超时 (timeout={timeout}s)")
            return None
        except requests.exceptions.ConnectionError:
            logger.error("工作流API连接失败")
            return None
        except Exception as e:
            logger.error(f"工作流API调用异常: {str(e)}")
            return None
    
    def _make_request(self, method: str, endpoint: str, data: Dict[str, Any] = None) -> Optional[Dict[str, Any]]:
        """发送HTTP请求"""
        if not self.enabled:
            logger.debug("工作流集成未启用，跳过调用")
            return None
            
        if not self.base_url:
            logger.warning("工作流基础URL未配置")
            return None
        
        url = f"{self.base_url}{endpoint}"
        
        try:
            logger.info(f"调用工作流API: {method} {url}")
            
            response = requests.request(
                method=method,
                url=url,
                headers=self.headers,
                json=data,
                timeout=self.timeout
            )
            
            # 检查响应状态
            if response.status_code in [200, 201]:
                result = response.json()
                logger.success(f"工作流API调用成功: {response.status_code}")
                return result
            else:
                logger.error(f"工作流API调用失败: {response.status_code} - {response.text}")
                return None
                
        except requests.exceptions.Timeout:
            logger.error(f"工作流API调用超时 (timeout={self.timeout}s)")
            return None
        except requests.exceptions.ConnectionError:
            logger.error("工作流API连接失败")
            return None
        except Exception as e:
            logger.error(f"工作流API调用异常: {str(e)}")
            return None
    
    def notify_knowledge_base_update(self, action: str, dataset_id: str, dataset_name: str, 
                                   document_info: Dict[str, Any] = None, 
                                   user_info: Dict[str, Any] = None) -> bool:
        """
        通知工作流知识库更新 - 根据知识库ID找到对应的工作流配置并发送通知
        
        Args:
            action: 操作类型 (create, update, delete, upload, etc.)
            dataset_id: 知识库ID
            dataset_name: 知识库名称
            document_info: 文档相关信息（可选）
            user_info: 用户信息（可选）
            
        Returns:
            bool: 是否成功发送通知
        """
        # 获取该知识库对应的工作流配置
        workflow_config = self._get_workflow_config_by_dataset(dataset_id)
        if not workflow_config:
            logger.debug(f"知识库 {dataset_id} 未配置工作流，跳过通知")
            return True  # 未配置工作流不视为错误
        
        payload = {
            "event": "knowledge_base_update",
            "action": action,
            "timestamp": self._get_timestamp(),
            "data": {
                "dataset": {
                    "id": dataset_id,
                    "name": dataset_name
                }
            }
        }
        
        # 添加文档信息（如果有）
        if document_info:
            payload["data"]["document"] = document_info
            
        # 添加用户信息（如果有）
        if user_info:
            payload["user"] = user_info
        else:
            # 默认用户信息
            payload["user"] = {
                "id": "system",
                "name": "Dify KB System"
            }
        
        # 添加元数据
        payload["metadata"] = {
            "source": "dify-kb-client",
            "version": "1.0.0",
            "environment": "production"
        }
        
        logger.info(f"发送知识库更新通知: {action} - {dataset_name} ({dataset_id})")
        
        result = self._make_request_to_workflow(workflow_config, payload)
        
        if result:
            logger.success("工作流通知发送成功")
            return True
        else:
            logger.warning("工作流通知发送失败，但知识库操作已完成")
            return False
    
    def notify_document_update(self, dataset_id: str, document_id: str, 
                               event_type: str, document_info: Dict[str, Any]) -> bool:
        """
        通知工作流文档更新
        
        Args:
            dataset_id: 知识库ID
            document_id: 文档ID
            event_type: 事件类型 (document_created, document_uploaded, document_updated, document_deleted)
            document_info: 文档信息
            
        Returns:
            bool: 是否通知成功
        """
        # 获取该知识库的配置信息
        kb_config = settings.get_knowledge_base_config(dataset_id)
        if not kb_config:
            logger.debug(f"知识库 {dataset_id} 未配置，跳过通知")
            return True  # 未配置不算错误
        
        workflow_config = kb_config.get("workflow_config")
        workflow_key = kb_config.get("workflow_key", "")
        
        if not workflow_config:
            logger.debug(f"知识库 {dataset_id} 未配置工作流，跳过通知")
            return True  # 未配置工作流不算错误
        
        # 构建通知数据
        notification_data = {
            "event": event_type,
            "timestamp": self._get_timestamp(),
            "workflow_key": workflow_key,  # 包含工作流key
            "dataset_info": {
                "dataset_id": dataset_id,
                "dataset_name": document_info.get("dataset_name", ""),
                "workflow_key": workflow_key
            },
            "document_info": {
                "document_id": document_id,
                "document_name": document_info.get("name", ""),
                "document_type": document_info.get("type", "text"),
                "size": document_info.get("size", 0),
                "word_count": document_info.get("word_count", 0),
                "indexing_status": document_info.get("indexing_status", "waiting"),
                "created_at": document_info.get("created_at", self._get_timestamp())
            }
        }
        
        # 发送通知
        return self._make_request_to_workflow(workflow_config, notification_data)
    
    def test_connection(self) -> bool:
        """测试工作流API连接"""
        if not self.enabled:
            logger.info("工作流集成未启用")
            return True
            
        if not self.base_url:
            logger.warning("工作流基础URL未配置")
            return False
        
        # 尝试调用健康检查接口（假设工作流API有health端点）
        health_endpoint = "/health"  # 可以根据实际情况调整
        result = self._make_request("GET", health_endpoint)
        
        if result:
            logger.success("工作流API连接测试成功")
            return True
        else:
            logger.error("工作流API连接测试失败")
            return False
    
    def _get_timestamp(self) -> str:
        """获取当前时间戳"""
        from datetime import datetime
        return datetime.now().isoformat()


# 创建全局工作流客户端实例
workflow_client = WorkflowClient()